import cv2
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision import transforms
import test_digit  # 用于模型预测的代码
from PIL import Image

img_size = 28  # 28x28是MNIST的图片训练集尺寸
kernel_connect = np.array([[1, 1, 1], [1, 0, 1], [1, 1, 1]], np.uint8)  # 膨胀化用的参数
ans = []  # 保存图片数组



# 载入训练好的 PyTorch 模型
import torch
import torch.nn as nn
import torch.nn.functional as F


class CNNModel(nn.Module):
    def __init__(self):
        super(CNNModel, self).__init__()

        # 卷积层 1
        self.conv1 = nn.Conv2d(1, 32, kernel_size=3, padding=1)
        # 卷积层 2
        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1)

        # 池化层（2x2的最大池化）
        self.pool = nn.MaxPool2d(2, 2)  # 这里定义了池化层

        # 全连接层
        self.fc1 = nn.Linear(64 * 7 * 7, 128)  # 64通道的7x7大小（假设输入为28x28图像）
        self.fc2 = nn.Linear(128, 10)  # 输出10个数字类别（0-9）

    def forward(self, x):
        x = self.conv1(x)  # 第一层卷积
        x = F.relu(x)
        x = self.pool(x)  # 池化操作

        x = self.conv2(x)  # 第二层卷积
        x = F.relu(x)
        x = self.pool(x)  # 池化操作

        x = x.view(x.size(0), -1)  # 展平
        x = self.fc1(x)  # 全连接层1
        x = self.fc2(x)  # 全连接层2
        return x


# 加载模型权重
model = CNNModel()
model.load_state_dict(torch.load('E:/mnist/mnist_model.pth'))
model.eval()  # 设置为评估模式

# 定义图像预处理转换
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5,))
])


def show(name):
    cv2.imshow("show", name)
    cv2.waitKey(0)


def split_digits(s, prefix_name):
    s = np.rot90(s)  # 使图片逆时针旋转90°
    s_copy = cv2.dilate(s, kernel_connect, iterations=1)  # 膨胀处理
    s_copy2 = s_copy.copy()

    contours, hierarchy = cv2.findContours(s_copy2, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)  # 检测轮廓

    idx = 0
    for contour in contours:
        idx = idx + 1
        [x, y, w, h] = cv2.boundingRect(contour)  # 获取轮廓的边界框
        digit = s_copy[y:y + h, x:x + w]

        # Padding调整宽高比，确保数字居中
        pad_len = (h - w) // 2
        if pad_len > 0:
            digit = cv2.copyMakeBorder(digit, 0, 0, pad_len, pad_len, cv2.BORDER_CONSTANT, value=0)
        elif pad_len < 0:
            digit = cv2.copyMakeBorder(digit, -pad_len, -pad_len, 0, 0, cv2.BORDER_CONSTANT, value=0)

        # 适当的空白边框
        pad = digit.shape[0] // 4  # 为了避免数字与边框相连，边框宽度设为4
        digit = cv2.copyMakeBorder(digit, pad, pad, pad, pad, cv2.BORDER_CONSTANT, value=0)

        # 缩放至28x28
        digit = cv2.resize(digit, (img_size, img_size), interpolation=cv2.INTER_AREA)
        digit = np.rot90(digit, 3)  # 逆时针旋转270°确保数字朝向正确

        # 将图像保存为文件
        cv2.imwrite(prefix_name + str(idx) + '.jpg', digit)
        ans.append(digit)

    # 进行数字识别
    test_digit.dj(ans, model, transform)


if __name__ == '__main__':
    img = cv2.imread('student_id.png')  # 输入的数字串图像路径
    img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)  # 转换为灰度图像
    ret, thresh_img = cv2.threshold(img, 127, 255, cv2.THRESH_BINARY_INV)  # 二值化处理，反转颜色
    show(thresh_img)  # 显示预处理后的图像
    split_digits(thresh_img, "split_img/split_img")  # 分割数字并保存

